#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <time.h>
#include <sys/time.h>

#define MAX_THREADS 16
#define M 1000000 // 1M

int *A;
long long total_sum = 0;
int n;
int num_threads;
pthread_mutex_t mutex;

// 线程参数结构体
typedef struct {
    int thread_id;
    long long partial_sum;
} ThreadParam;

// 数组求和线程函数
void *array_sum(void *arg) {
    ThreadParam *param = (ThreadParam *)arg;
    int id = param->thread_id;
    
    // 计算每个线程负责的范围
    int elements_per_thread = n / num_threads;
    int start = id * elements_per_thread;
    int end = (id == num_threads - 1) ? n : start + elements_per_thread;
    
    long long sum = 0;
    for (int i = start; i < end; i++) {
        sum += A[i];
    }
    
    param->partial_sum = sum;
    
    // 使用互斥锁保护全局变量
    pthread_mutex_lock(&mutex);
    total_sum += sum;
    pthread_mutex_unlock(&mutex);
    
    pthread_exit(NULL);
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        printf("Usage: %s <n (in millions)> <num_threads>\n", argv[0]);
        return 1;
    }
    
    n = atoi(argv[1]) * M;
    num_threads = atoi(argv[2]);
    
    if (num_threads < 1 || num_threads > MAX_THREADS) {
        printf("Number of threads must be between 1 and %d\n", MAX_THREADS);
        return 1;
    }
    
    srand(time(NULL));
    
    // 初始化数组
    A = (int *)malloc(n * sizeof(int));
    for (int i = 0; i < n; i++) {
        A[i] = rand() % 100;
    }
    
    // 初始化互斥锁
    pthread_mutex_init(&mutex, NULL);
    
    // 创建线程
    pthread_t threads[MAX_THREADS];
    ThreadParam params[MAX_THREADS];
    
    struct timeval start, end;
    gettimeofday(&start, NULL);
    
    for (int i = 0; i < num_threads; i++) {
        params[i].thread_id = i;
        params[i].partial_sum = 0;
        pthread_create(&threads[i], NULL, array_sum, &params[i]);
    }
    
    // 等待所有线程完成
    for (int i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }
    
    gettimeofday(&end, NULL);
    double time_taken = (end.tv_sec - start.tv_sec) * 1e6;
    time_taken = (time_taken + (end.tv_usec - start.tv_usec)) * 1e-6;
    
    // 输出结果
    printf("Array (first 10 elements): ");
    for (int i = 0; i < 10 && i < n; i++) {
        printf("%d ", A[i]);
    }
    printf("\n...\n");
    
    printf("Total sum: %lld\n", total_sum);
    printf("Time taken: %f seconds\n", time_taken);
    
    // 清理
    free(A);
    pthread_mutex_destroy(&mutex);
    
    return 0;
}